import os
import xml.etree.ElementTree as ET


names =  ["xishuai","xiewenyee","zangjia","duncifuyajia","daozongjuanyeming","hongchun","zhonghualanbujia","lougu",
           "miyee","tiancaiyee","lanlvqijiaoming","nianchong","duchong","youmeitaie","xiaocaie","jinguizi","caoling","daolvchun",
           "xingbaixuedenge","maimuyeming","dilaohu","songmaochong","chenwudenge","xiusefuzichun","tutaie","huangzuliechun","xiumuyee",
           "woweiyee","taozhuming","huangdue","baiheihuataie","baitiaoyee","erhuaming","mianlingchong","yinwenyee","longshi"] # 改成自己的类别

pest24_name = ['Bollworm', 'Meadow borer', 'Gryllotalpa orientalis', 'Little Gecko', 'Agriotes fuscicollis Miwa', 'Nematode trench',
        'Athetis lepigone', 'Scotogramma trifolii Rottemberg', 'Armyworm', 'Spodoptera cabbage', 'Anomala corpulenta',
        'Spodoptera exigua', 'Plutella xylostella', 'holotrichia parallela', 'Rice planthopper', 'Yellow tiger',
        'Land tiger', 'eight-character tiger', 'holotrichia oblita', 'Stem borer', 'Striped rice bore',
        'Rice Leaf Roller', 'Spodoptera litura', 'Melahotus']


# D:/rice_dataset_2/3_稻飞虱/labels
class_dic = {}
# class_dic['lougu'] = 'Gryllotalpa orientalis'
# class_dic['daozongjuanyeming'] = 'Rice Leaf Roller'
# class_dic['xiewenyee'] = 'Spodoptera litura'
# class_dic['nianchong'] = 'Armyworm'
# class_dic['xiaocaie'] = 'Plutella xylostella'
# class_dic['mianlingchong'] = 'Bollworm'
# class_dic['erhuaming'] = 'Striped rice borer'
# class_dic['tiancaiyee'] = 'Spodoptera exigua'

class_dic['Gryllotalpa orientalis'] = 'lougu'
class_dic['Rice Leaf Roller'] = 'daozongjuanyeming'
class_dic['Spodoptera litura'] = 'xiewenyee'
class_dic['Armyworm'] = 'nianchong'
class_dic['Plutella xylostella'] = 'xiaocaie'
class_dic['Bollworm'] = 'mianlingchong'
class_dic['Striped rice bore'] = 'erhuaming'
class_dic['Spodoptera exigua'] = 'tiancaiyee'

xml_fold = "E:/liyuhai_dir/dataset/pest24/Annotations"

def changeName(xml_fold):

    '''
    xml_fold: xml存放文件夹
    origin_name: 原始名字，比如弄错的名字，原先要cow,不小心打成cwo
    new_name: 需要改成的正确的名字，在上个例子中就是cow
    '''
    files = os.listdir(xml_fold)
    cnt = 0
    for xmlFile in files:
        file_path = os.path.join(xml_fold, xmlFile)
        tree = ET.parse(file_path)
        root = tree.getroot()
        for obj in root.findall('object'):#获取object节点中的name子节点
            tmp_name = obj.find('name').text
            if tmp_name in class_dic:
                obj.find('name').text = class_dic[tmp_name]
                print("change %s to %s." % (tmp_name, class_dic[tmp_name]))
                cnt += 1
            else:
                root.remove(obj)   # 删除这个数据
                cnt += 1
        tree.write(file_path, xml_declaration=True) # 保存到指定文件
    print("有%d个文件被成功修改。" % cnt)

changeName(xml_fold )
